import os
import sys
import io
from pathlib import Path
from typing import Optional, Tuple

import fitz  # PyMuPDF
from PIL import Image
import json
from PIL import ImageFilter, ImageOps
from typing import Literal
try:
    from rembg import remove as rembg_remove  # type: ignore
    HAS_REMBG = True
except Exception:
    HAS_REMBG = False


TARGET_KB = 200
MIN_QUALITY = 35
QUALITY_STEP = 5
INITIAL_QUALITY = 85
MAX_DOWNSCALE_STEPS = 4
DOWNSCALE_FACTOR = 0.85  # each step reduces dimensions by 15%
RESULT_DIR_NAME = "result"

# 常见证件照尺寸（毫米）
# 注：不同机构要求可能略有差异，提供通用版本，并支持自定义
PRESET_MM: dict[str, tuple[int, int]] = {
    "一寸证件照": (25, 35),
    "两寸证件照": (35, 49),
    "护照照片": (33, 48),
    "大一寸": (33, 48),
    "小一寸": (22, 32),
    "英语四六级": (35, 49),
    "教师资格证": (35, 45),
    "国家公务员考试": (25, 35),
    "导游资格考试": (25, 35),
    "简历": (25, 35),
    "学籍照片": (35, 45),
    "网上报名": (35, 45),
}


def ensure_jpeg_under_size(img: Image.Image, out_path: Path, target_kb: int = TARGET_KB,
                           initial_quality: int = INITIAL_QUALITY,
                           min_quality: int = MIN_QUALITY,
                           quality_step: int = QUALITY_STEP,
                           allow_downscale: bool = True) -> bool:
    """Save image as JPEG under target size. Iteratively decrease quality and optionally downscale.

    Returns True if successful, False otherwise.
    """
    img = img.convert("RGB")  # ensure no alpha for JPEG

    def try_save(current_img: Image.Image, quality: int) -> int:
        current_img.save(out_path, format="JPEG", quality=quality, optimize=True, progressive=True)
        return out_path.stat().st_size

    quality = initial_quality
    # First pass: only adjust quality
    while quality >= min_quality:
        size = try_save(img, quality)
        if size <= target_kb * 1024:
            return True
        quality -= quality_step

    # Optional second pass: downscale if still too large
    if allow_downscale:
        w, h = img.size
        temp_img = img
        for _ in range(MAX_DOWNSCALE_STEPS):
            w = int(w * DOWNSCALE_FACTOR)
            h = int(h * DOWNSCALE_FACTOR)
            if w < 64 or h < 64:
                break
            temp_img = temp_img.resize((w, h), Image.LANCZOS)
            quality = min_quality  # keep minimum quality to prioritize size
            size = try_save(temp_img, quality)
            if size <= target_kb * 1024:
                return True

    return False


def ensure_png_under_size(img: Image.Image, out_path: Path, target_kb: int = TARGET_KB) -> bool:
    """Save image as PNG under target size by downscaling progressively.
    PNG是无损格式，主要通过降采样与最高压缩级别实现体积控制。
    """
    # Keep mode as RGB to avoid palette issues for photos
    if img.mode not in ("RGB", "RGBA"):
        img = img.convert("RGB")

    def try_save(current_img: Image.Image) -> int:
        current_img.save(out_path, format="PNG", optimize=True, compress_level=9)
        return out_path.stat().st_size

    # initial attempt
    size = try_save(img)
    if size <= target_kb * 1024:
        return True

    # downscale loop
    w, h = img.size
    temp_img = img
    for _ in range(MAX_DOWNSCALE_STEPS):
        w = int(w * DOWNSCALE_FACTOR)
        h = int(h * DOWNSCALE_FACTOR)
        if w < 64 or h < 64:
            break
        temp_img = temp_img.resize((w, h), Image.LANCZOS)
        size = try_save(temp_img)
        if size <= target_kb * 1024:
            return True
    return False


def ensure_pdf_under_size(img: Image.Image, out_path: Path, target_kb: int = TARGET_KB) -> bool:
    """Embed image into a single-page PDF aiming for target size.
    通过逐步降采样控制体积；PDF额外开销不可控，因此仅做尽力控制。
    """
    if img.mode != "RGB":
        img = img.convert("RGB")

    def try_save(current_img: Image.Image) -> int:
        current_img.save(out_path, format="PDF")
        return out_path.stat().st_size

    size = try_save(img)
    if size <= target_kb * 1024:
        return True

    w, h = img.size
    temp_img = img
    for _ in range(MAX_DOWNSCALE_STEPS):
        w = int(w * DOWNSCALE_FACTOR)
        h = int(h * DOWNSCALE_FACTOR)
        if w < 64 or h < 64:
            break
        temp_img = temp_img.resize((w, h), Image.LANCZOS)
        size = try_save(temp_img)
        if size <= target_kb * 1024:
            return True
    return False


def save_image_with_format(img: Image.Image, out_path: Path, output_format: str, target_kb: int) -> bool:
    fmt = output_format.lower()
    if fmt in ("jpg", "jpeg"):
        return ensure_jpeg_under_size(img, out_path, target_kb)
    elif fmt == "png":
        return ensure_png_under_size(img, out_path, target_kb)
    elif fmt == "pdf":
        return ensure_pdf_under_size(img, out_path, target_kb)
    else:
        # default to jpg
        return ensure_jpeg_under_size(img, out_path, target_kb)


def pdf_page_to_image(doc: fitz.Document, page_index: int, dpi: int = 200) -> Image.Image:
    """Render a PDF page to a PIL Image at given DPI."""
    page = doc.load_page(page_index)
    zoom = dpi / 72.0
    mat = fitz.Matrix(zoom, zoom)
    pix = page.get_pixmap(matrix=mat, alpha=False)
    mode = "RGB" if pix.n < 4 else "RGBA"
    img = Image.frombytes(mode, [pix.width, pix.height], pix.samples)
    if img.mode != "RGB":
        img = img.convert("RGB")
    return img


def convert_pdf_to_outputs(pdf_path: Path, root: Path, result_root: Path, target_kb: int, output_format: str, status_file: Optional[Path] = None) -> None:
    """Convert each page of a PDF to selected format, aiming for target_kb.

    Output filename pattern: <basename>_p<page> .<ext>
    """
    try:
        doc = fitz.open(pdf_path)
    except Exception as e:
        print(f"[PDF ERROR] 无法打开: {pdf_path} -> {e}")
        return

    base = pdf_path.stem
    out_dir = result_root / pdf_path.relative_to(root).parent
    out_dir.mkdir(parents=True, exist_ok=True)
    for i in range(doc.page_count):
        ext = output_format.lower()
        out_name = f"{base}_p{i+1}.{ext}"
        out_path = out_dir / out_name
        try:
            if status_file:
                try:
                    status_file.write_text(json.dumps({"phase": "processing", "filename": f"{pdf_path.name} 第{i+1}页"}, ensure_ascii=False), encoding="utf-8")
                except Exception:
                    pass
            img = pdf_page_to_image(doc, i, dpi=200)
            ok = save_image_with_format(img, out_path, output_format, target_kb)
            if ok:
                print(f"[PDF->IMG] {pdf_path.name} 第{i+1}页 -> {out_path} ({out_path.stat().st_size/1024:.0f}KB)")
            else:
                print(f"[WARN] {pdf_path.name} 第{i+1}页无法压到 {target_kb}KB，已保存但可能超限 ({out_path.stat().st_size/1024:.0f}KB)")
        except Exception as e:
            print(f"[PAGE ERROR] {pdf_path.name} 第{i+1}页转换失败: {e}")


def process_image_file(file_path: Path, root: Path, result_root: Path, target_kb: int, output_format: str, status_file: Optional[Path] = None) -> None:
    """Process existing image and write to result folder in selected format, aiming for target_kb."""
    try:
        with Image.open(file_path) as img:
            if status_file:
                try:
                    status_file.write_text(json.dumps({"phase": "processing", "filename": file_path.name}, ensure_ascii=False), encoding="utf-8")
                except Exception:
                    pass
            out_dir = result_root / file_path.relative_to(root).parent
            out_dir.mkdir(parents=True, exist_ok=True)
            ext = output_format.lower()
            out_path = out_dir / (file_path.stem + f".{ext}")
            # Try to keep dimensions; downscale only if necessary
            ok = save_image_with_format(img, out_path, output_format, target_kb)
            size_kb = out_path.stat().st_size / 1024 if out_path.exists() else 0
            if ok:
                print(f"[IMG处理] {file_path} -> {out_path} ({size_kb:.0f}KB)")
            else:
                print(f"[WARN] 无法将 {file_path} 压至 {target_kb}KB，当前 {size_kb:.0f}KB")
    except Exception as e:
        print(f"[IMG ERROR] 无法处理: {file_path} -> {e}")


def scan_and_process(root: Path, target_kb: int = TARGET_KB, output_format: str = "jpg", status_file: Optional[Path] = None) -> Tuple[int, int, int]:
    result_root = root / RESULT_DIR_NAME
    result_root.mkdir(parents=True, exist_ok=True)
    pdf_count = 0
    img_count = 0
    out_count = 0
    for path in root.rglob("*"):
        if path.is_file():
            # 跳过空文件，避免后续 PDF/IMG 读取报错
            try:
                if path.stat().st_size == 0:
                    print(f"[SKIP] 空文件：{path}")
                    continue
            except Exception:
                pass
            suffix = path.suffix.lower()
            # 跳过 result 目录中的内容，避免重复处理
            try:
                if path.is_relative_to(result_root):
                    continue
            except AttributeError:
                if str(path).startswith(str(result_root)):
                    continue
            if suffix == ".pdf":
                pdf_count += 1
                convert_pdf_to_outputs(path, root, result_root, target_kb, output_format, status_file=status_file)
                out_count += 1
            elif suffix in {".jpg", ".jpeg", ".png"}:
                img_count += 1
                process_image_file(path, root, result_root, target_kb, output_format, status_file=status_file)
                out_count += 1
    print(f"\n完成：处理PDF {pdf_count} 个，图片 {img_count} 个，总输出 {out_count} 个。")
    return pdf_count, img_count, out_count


def main():
    if len(sys.argv) > 1:
        root = Path(sys.argv[1]).resolve()
    else:
        root = Path.cwd()
    # Optional args: target_kb and output_format
    target_kb = TARGET_KB
    output_format = "jpg"
    if len(sys.argv) > 2:
        try:
            target_kb = int(sys.argv[2])
        except Exception:
            pass
    if len(sys.argv) > 3:
        output_format = sys.argv[3]
    print(f"扫描目录：{root}，目标大小：{target_kb}KB，输出格式：{output_format}")
    scan_and_process(root, target_kb, output_format)


def recolor_id_photo(input_path: Path, bg_color: Literal["blue","red","white"], output_path: Path,
                     threshold: int = 40) -> bool:
    """换底色前先进行人物抠图，避免面部被改色。

    - 优先使用 rembg(U2Net) 人像分割得到透明背景的人像，然后与纯色背景合成。
    - 若环境未安装 rembg，回退到边缘采样背景估计的启发式方法。
    """
    target_map = {
        "blue": (70, 130, 220),
        "red": (210, 70, 70),
        "white": (255, 255, 255),
    }
    tgt = target_map.get(bg_color, (255, 255, 255))
    try:
        # 1) 优选：rembg 人像抠图
        if HAS_REMBG:
            try:
                with open(input_path, "rb") as f:
                    input_bytes = f.read()
                output_bytes = rembg_remove(input_bytes)
                fg = Image.open(io.BytesIO(output_bytes)).convert("RGBA")
                bg = Image.new("RGBA", fg.size, (*tgt, 255))
                # 叠加人像到新背景，保持边缘过渡
                alpha = fg.split()[3]
                bg.paste(fg, (0, 0), alpha)
                bg.convert("RGB").save(output_path, format="PNG")
                return True
            except Exception as e:
                print(f"[REMBG FALLBACK] {input_path} -> {e}")
                # 继续走回退方案

        # 2) 回退：边界采样估计背景并替换
        with Image.open(input_path) as img:
            if img.mode != "RGB":
                img = img.convert("RGB")
            w, h = img.size
            border = 4
            pixels = []
            px = img.load()
            for x in range(w):
                for y in range(border):
                    pixels.append(px[x, y])  # top
                for y in range(h - border, h):
                    pixels.append(px[x, y])  # bottom
            for y in range(h):
                for x in range(border):
                    pixels.append(px[x, y])  # left
                for x in range(w - border, w):
                    pixels.append(px[x, y])  # right
            if not pixels:
                return False
            r = sum(p[0] for p in pixels) / len(pixels)
            g = sum(p[1] for p in pixels) / len(pixels)
            b = sum(p[2] for p in pixels) / len(pixels)
            br, bgc, bb = int(r), int(g), int(b)

            mask = Image.new("L", (w, h), 0)
            m = mask.load()
            for y in range(h):
                for x in range(w):
                    pr, pg, pb = px[x, y]
                    dr = pr - br
                    dg = pg - bgc
                    db = pb - bb
                    d = (dr * dr + dg * dg + db * db) ** 0.5
                    m[x, y] = 255 if d <= threshold else 0

            mask = mask.filter(ImageFilter.GaussianBlur(radius=2))
            bg_img = Image.new("RGB", (w, h), tgt)
            out = Image.composite(bg_img, img, mask)
            out.save(output_path)
            return True
    except Exception as e:
        print(f"[RECOLOR ERROR] {input_path} -> {e}")
        return False


def _mm_to_px(mm: int | float, dpi: int) -> int:
    return max(1, int(round((mm / 25.4) * dpi)))


def resize_id_photo(input_path: Path, output_path: Path, *,
                    preset_name: str | None = None,
                    width_mm: int | float | None = None,
                    height_mm: int | float | None = None,
                    dpi: int = 300,
                    fit: Literal["cover", "contain"] = "cover",
                    bg_color: Literal["white", "blue", "red"] = "white") -> bool:
    """将证件照调整到指定尺寸（支持预设与自定义）。

    - preset_name：使用 PRESET_MM 里的常见尺寸；如传入 "自定义" 或 None，则使用 width_mm/height_mm。
    - width_mm/height_mm：自定义毫米尺寸；至少需要其中两个及 dpi 才能计算目标像素。
    - dpi：用于像素换算，默认 300。
    - fit：cover（充满裁剪）或 contain（完整适配留边）。
    - bg_color：当使用 contain 时用于填充留边；同时也作为最终输出的纯色背景（PNG/JPG 都可）。
    """
    try:
        with Image.open(input_path) as img:
            if img.mode not in ("RGB", "RGBA"):
                img = img.convert("RGB")

            # 目标毫米尺寸
            if preset_name and preset_name in PRESET_MM:
                wmm, hmm = PRESET_MM[preset_name]
            else:
                wmm = int(width_mm or 25)
                hmm = int(height_mm or 35)

            tw = _mm_to_px(wmm, dpi)
            th = _mm_to_px(hmm, dpi)
            target_ar = tw / th

            # cover：中心裁剪到目标比例
            if fit == "cover":
                w0, h0 = img.size
                src_ar = w0 / float(h0)
                if src_ar > target_ar:
                    # 宽太宽，裁掉左右
                    new_w = int(round(h0 * target_ar))
                    x0 = (w0 - new_w) // 2
                    crop_box = (x0, 0, x0 + new_w, h0)
                else:
                    # 高太高，裁掉上下
                    new_h = int(round(w0 / target_ar))
                    y0 = (h0 - new_h) // 2
                    crop_box = (0, y0, w0, y0 + new_h)
                img = img.crop(crop_box)
                img = img.resize((tw, th), Image.LANCZOS)
                img.save(output_path)
                return True

            # contain：完整适配，居中留边
            else:
                w0, h0 = img.size
                r = min(tw / float(w0), th / float(h0))
                new_size = (max(1, int(round(w0 * r))), max(1, int(round(h0 * r))))
                resized = img.resize(new_size, Image.LANCZOS)
                # 背景填充
                bg_map = {
                    "white": (255, 255, 255),
                    "blue": (70, 130, 220),
                    "red": (210, 70, 70),
                }
                bgr = bg_map.get(bg_color, (255, 255, 255))
                canvas = Image.new("RGB", (tw, th), bgr)
                cx = (tw - resized.size[0]) // 2
                cy = (th - resized.size[1]) // 2
                canvas.paste(resized, (cx, cy))
                canvas.save(output_path)
                return True
    except Exception as e:
        print(f"[RESIZE ERROR] {input_path} -> {e}")
        return False


if __name__ == "__main__":
    main()